#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/DTI.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore//DWI.h"

#include <teem/ell.h>
#include <teem/ten.h>

#include <omp.h>

#include <levmar/levmar.h>

#define _USE_MATH_DEFINES
#include <math.h>
#include <time.h>

#include <string>

using namespace std;

char *pDWIPathname = NULL;				/* input DWI pathname */
char *pDTIPathname = NULL;				/* input diffusion tensor pathname */
char *pWMMASKPathname = NULL;			/* input white matter mask pathname */
char *pModelSelectionPathname = NULL;	/* input model selection pathname */
char *pFiberPathname = NULL;			/* output fiber configure pathname */
float stepSize = 0.2f;
float maxLength = 5.0f;
float faThreshold = 0.001f;
float angleThreshold = 0.707f;
int scale = 1;


/* MLS fiber tracking parameters */
float MLSSize = 4.0f;
CPoint3F *pSamplePoints = NULL;
float *pSamplePointsWeight = NULL;

CDWI *pDWI = NULL;					/* DWI */
CDTI *pDTI = NULL;					/* DTI */
CScalar *pWMMASK = NULL;			/* white matter */
CScalar *pModelSelection = NULL;	/* model selection */

struct AData {
	float s0;
	float bValue;
	CVector3F *gradientOrientations;
	float lam3;
};

void PrintUsage()
{
	printf("Usage: ");
	printf("DTIFiberTracking -dti <DTI> -dwi <DWI> -wm_mask <White Matter MASK> -fiber <Fiber> [-stepsize <stepsize> -maxlength <maxlength> -fa_thresold <fa_threshold> -angle_threshold <fa_threshold> -scale <scale>]\n");
	printf("\n");
	//printf("-vec <VectorField> = input VectorField nrrd file\n");
	//printf("-fiber <Fiber> = output Fiber configure file\n");
	//printf("-stepsize <stepsize> = step size\n");
	//printf("-maxlength <maxlength> = maximum fiber length\n");
	//printf("-fa_threshold <fa_thresold> = FA threshold used to differ white matter and grey matter\n");
	//printf("-angle_threshold <angle_thresold> = angle threshold between two directions, in cosine\n");
	//printf("-scale <scale> = scale of the result\n");
}


/* Form a 2D tensor from the parameters */
void PeledRotate2D(double ten[7], double lam1, double lam3, double phi) {
  double cc, ss, d3, d1, d2;

  cc = cos(phi);
  ss = sin(phi);
  d1 = cc*cc*lam1 + ss*ss*lam3;
  d3 = cc*ss*(lam1 - lam3);
  d2 = ss*ss*lam1 + cc*cc*lam3;

  TEN_T_SET(ten, 1.0, d1, d3, 0.0, d2, 0.0, lam3);
  return;
}

/* call back function for levmar optimaization */
void TwoTensorFittingCB(double *pp, double *xx, int mm, int nn, void *_pvlData) {
  /* static const char me[]="_tenLevmarPeledCB"; */
  double tenA[7], tenB[7];
  int ii;
  AData *pvlData;
  CVector3F *gradients;
 
  AIR_UNUSED(mm);
  pvlData = AIR_CAST(AData *, _pvlData);

  /* Form the tensors using the estimated parms */
  PeledRotate2D(tenA, pp[0], pvlData->lam3, pp[2]);
  PeledRotate2D(tenB, pp[0], pvlData->lam3, pp[3]);

  gradients = AIR_CAST(CVector3F *, pvlData->gradientOrientations);

  for (ii = 0; ii < nn; ii++) {
    double argA, argB, sigA, sigB;
	double egrad[3];
	egrad[0] = gradients[ii].m_x;
	egrad[1] = gradients[ii].m_y;
	egrad[2] = gradients[ii].m_z;
	argA = -pvlData->bValue * TEN_T3V_CONTR(tenA, egrad);
	argB = -pvlData->bValue * TEN_T3V_CONTR(tenB, egrad);
	sigA = exp(argA);
	sigB = exp(argB);

    xx[ii] = pp[1] * sigA + (1.0 - pp[1]) * sigB;
  }
  return;
}

void Fitting(float *_dwi, const float s0, const float bValue, 
	float *eval, float *evec, CVector3F *gradients, const int n, float *twoTensor)
{
	double guess[4], loBnd[4], upBnd[4], opts[LM_OPTS_SZ], info[LM_INFO_SZ];
	double tenA[7], tenB[7];
	double matA[9], matB[9], rott[9], matTemp[9];
	double *dwi = new double[n];

	// dwi
	for (int i = 0; i < n; ++i) {
		dwi[i] = _dwi[i] / s0;
	}

	// set lower and upper bound
	loBnd[0] = 0.0;
	loBnd[1] = 0.0;
	loBnd[2] = -M_PI / 2.0;
	loBnd[3] = -M_PI / 2.0;
	upBnd[0] = eval[0] * 5.0;
	upBnd[1] = 1.0;
	upBnd[2] = M_PI / 2.0;
	upBnd[3] = M_PI / 2.0;

	// set initial guess
	guess[0] = eval[0];
	guess[1] = 0.5;
	guess[2] = M_PI / 4.0;
	guess[3] = -M_PI / 4.0;

	// set constraints for LM optimization
	opts[0] = 1E-3;
	opts[1] = 1E-8;
	opts[2] = 1E-8;
	opts[3] = 1E-8;
	opts[4] = 1E-8;

	AData aData;
	aData.s0 = s0;
	aData.bValue = bValue;
	aData.gradientOrientations = gradients;
	aData.lam3 = eval[2];

	int lmret = dlevmar_bc_dif(TwoTensorFittingCB, guess, dwi, 4, n, 
		loBnd, upBnd, NULL, 1000, opts, info, NULL, NULL, (void *)&aData);

	SafeDeleteArray(dwi);

	if (-1 == lmret) {
		printf("dlevmar_bc_dif() failed!\n");
	} else {
		// get two tensor fit result
		PeledRotate2D(tenA, guess[0], eval[2], guess[2]);
		PeledRotate2D(tenB, guess[0], eval[2], guess[3]);
		TEN_T2M(matA, tenA);
		TEN_T2M(matB, tenB);

		ELL_3M_TRANSPOSE(rott, evec);
		ELL_3M_MUL(matTemp, matA, evec);
		ELL_3M_MUL(matA, rott, matTemp);
		ELL_3M_MUL(matTemp, matB, evec);
		ELL_3M_MUL(matB, rott, matTemp);

		TEN_M2T(tenA, matA);
		TEN_M2T(tenB, matB);

		if (guess[1] > 0.5) {
			twoTensor[0] = 1.0f;
			twoTensor[7] = 1.0f - guess[1];

			for (int i = 1; i < 7; ++i) {
				twoTensor[i] = tenA[i];
				twoTensor[i+7] = tenB[i];
			}

		} else {
			twoTensor[0] = 1.0f;
			twoTensor[7] = guess[1];

			for (int i = 1; i < 7; ++i) {
				twoTensor[i] = tenB[i];
				twoTensor[i+7] = tenA[i];
			}
		}
		
	}
}

void TwoTensorFitting(float *dti, float *dwi, float s0, CVector3F *dirs)
{
	uint n = pDWI->GetN();
	uint m = pDWI->GetM();
	float bValue = pDWI->GetBValue();
	CVector3F *diffusionGradients = pDWI->GetDiffusionGradients();

	// compute eval and evec of DTI
	float eval[3], evec[9];
	pDTI->DTI2Eigen(dti, eval, evec);

	// rotate the graidents
	CVector3F *gradients = new CVector3F[n-m];
	for (uint k = m; k < n; ++k) {
		float v[3], vv[3];
		v[0] = diffusionGradients[k].m_x;
		v[1] = diffusionGradients[k].m_y;
		v[2] = diffusionGradients[k].m_z;
		ELL_3MV_MUL(vv, evec, v);
		gradients[k-m].m_x = vv[0];
		gradients[k-m].m_y = vv[1];
		gradients[k-m].m_z = vv[2];
	}

	// do the two tensor fitting
	float res[14];
	Fitting(dwi, s0, bValue, eval, evec, gradients, n-m, res);

	SafeDeleteArray(gradients);

	// tensor A
	pDTI->DTI2Direction(&(res[0]), &(dirs[0]));
	// tensor B
	pDTI->DTI2Direction(&(res[7]), &(dirs[1]));
}


int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dti") == 0) {
			pDTIPathname = argv[++i];
		} else if (strcmp(argv[i], "-dwi") == 0) {
			pDWIPathname = argv[++i];
		} else if (strcmp(argv[i], "-wm_mask") == 0) {
			pWMMASKPathname = argv[++i];
		} else if (strcmp(argv[i], "-fiber") == 0) {
			pFiberPathname = argv[++i];
		} else if (strcmp(argv[i], "-stepsize") == 0) {
			stepSize = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-maxlength") == 0) {
			maxLength = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-fa_threshold") == 0) {
			faThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-angle_threshold") == 0) {
			angleThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-scale") == 0) {
			scale = (uint)atoi(argv[++i]);
		} else {
            printf("unknown %s\n", argv[i]);
			return 1;		// error
		}
	}

	if (pDTIPathname == NULL || pWMMASKPathname == NULL || pFiberPathname == NULL)
		return -2;

	return 0;
}


void CreateMLSSamplePoints()
{
	SafeDeleteArray(pSamplePoints);
	SafeDeleteArray(pSamplePointsWeight);
	pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	pSamplePoints[0].m_x = 0.0f;
	pSamplePoints[0].m_y = 0.0f;
	pSamplePoints[0].m_z = 0.0f;
	pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(pSamplePoints[j], pt);
				if (dis < 0.25f) {
					flag = false;
					break;
				}
			}
		}

		pSamplePoints[i] = pt;
		pSamplePointsWeight[i] = InnerProduct(pSamplePoints[i], pSamplePoints[i]) * 4.0f;
		pSamplePointsWeight[i] = expf(pSamplePointsWeight[i] / -2.5f);
	}
}

void DTIMLSFilter(const CPoint3F pos, float *lastDTI, float *filteredDTI)
{
	float m[9];
	TEN_T2M(m, lastDTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	memset(filteredDTI, 0, sizeof(float)*7);
	float temp_DTI[7];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pWMMASK->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		pDTI->GetDTI(samples[i], temp_DTI);

		for (uint j = 0; j < 7; ++j) {
			filteredDTI[j] += pSamplePointsWeight[i] * temp_DTI[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < 7; ++j) {
		filteredDTI[j] = filteredDTI[j] / weight;
	}
}

void DWIMLSFilter(const CPoint3F pos, float *filteredDTI, float *filteredDWI, float *filteredS0)
{
	float m[9];
	TEN_T2M(m, filteredDTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	uint n = pDWI->GetN() - pDWI->GetM();
	memset(filteredDWI, 0, sizeof(float)*n);
	*filteredS0 = 0.0f;
	float *temp_DWI = new float[n];
	float temp_S0 = 0.0f;
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pWMMASK->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		pDWI->GetNS0(samples[i], temp_DWI);
		temp_S0 = pDWI->GetS0(samples[i]);

		for (uint j = 0; j < n; ++j) {
			filteredDWI[j] += pSamplePointsWeight[i] * temp_DWI[j];
		}
		*filteredS0 += pSamplePointsWeight[i] * temp_S0;
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < n; ++j) {
		filteredDWI[j] = filteredDWI[j] / weight;
	}
	*filteredS0 = *filteredS0 / weight;

	SafeDeleteArray(temp_DWI);
}

int ModelSelection(const CPoint3F pos)
{
	float model = pWMMASK->GetScalar(pos);
	if (model < 0.5f)
		return 0;
	else if (model < 1.5f)
		return 1;
	else
		return 2;
}


int GetNextDirection(const CPoint3F pos, CVector3F *in, float *lastDTI, CVector3F *out)
{
	float filteredDTI[7];
	DTIMLSFilter(pos, lastDTI, filteredDTI);
	memcpy(lastDTI, filteredDTI, sizeof(float)*7);

	uint n = pDWI->GetN() - pDWI->GetM();
	float *filteredDWI = new float[n];
	float filteredS0 = 0.0f;
	DWIMLSFilter(pos, filteredDTI, filteredDWI, &filteredS0);

	CVector3F dirs[2];
	int count = 0;
	int model = ModelSelection(pos);
	if (model == 0) {
		dirs[0] = CVector3F(0, 0, 0);
		count = 1;
	} else if (model == 1) {
		pDTI->DTI2Direction(filteredDTI, dirs);		
		count = 1;
	} else {
		/* two tensor fitting */
		TwoTensorFitting(filteredDTI, filteredDWI, filteredS0, dirs);
		count = 2;
	}

	SafeDeleteArray(filteredDWI);

	if (in == NULL) {
		for (int i = 0; i < count; ++i) {
			out[i] = dirs[i];
		}
		return count;
	} else {
		float max_v = InnerProduct(in, &(dirs[0]));
		int max_index = 0;
		for (int i = 1; i < count; ++i) {
			float v = InnerProduct(in, &(dirs[i]));
			if (fabs(v) > fabs(max_v)) {
				max_v = v;
				max_index = i;
			}
		}

		if (fabs(max_v) > angleThreshold) {
			if (max_v < 0.0f) {
				out->m_x = -dirs[max_index].m_x;
				out->m_y = -dirs[max_index].m_y;
				out->m_z = -dirs[max_index].m_z;
			} else {
				out->m_x = dirs[max_index].m_x;
				out->m_y = dirs[max_index].m_y;
				out->m_z = dirs[max_index].m_z;
			}
			return 1;
		} else {
			return 0;
		}
	}
}

int TrackingFibers(CFiber **fibers, CPoint3F pos)
{
	float initDTI[7] = {1, 1, 0, 0, 1, 0, 1};
	CVector3F dirs[3];
	int count = GetNextDirection(pos, NULL, initDTI, dirs);

	*fibers = new CFiber[count];
	
	for (int i = 0; i < count; ++i) {
		Seed seed;
		seed.pos = pos;
		seed.dir = dirs[i];

		int innerSteps = 5;
		uint maxSteps = (uint)(maxLength / (stepSize * (float)innerSteps));
		uint fnum, bnum;
		CPoint3F *pF = new CPoint3F[maxSteps];
		CPoint3F *pB = new CPoint3F[maxSteps];
		CPoint3F lastPos;
		CVector3F lastDir;
		float lastDTI[7];

		/* tracking forward */
		memcpy(lastDTI, initDTI, sizeof(float)*7);
		lastPos = seed.pos;
		lastDir = seed.dir;
		pF[0] = lastPos;
		for (fnum = 1; fnum < maxSteps; ++fnum) {
			CPoint3F nextPos;
			CVector3F nextDir;

			bool flag = true;
			for (int step = 0; step < innerSteps; ++step) {
				nextPos = lastPos + lastDir * stepSize;

				if (GetNextDirection(nextPos, &lastDir, lastDTI, &nextDir) == 0) {
					flag = false;
					break;
				}
				lastDir = nextDir;
				lastPos = nextPos;
			}
			if (flag == false)
				break;

			pF[fnum] = nextPos;
			lastDir = nextDir;
			lastPos = nextPos;
		}

		/* tracking backward */
		memcpy(lastDTI, initDTI, sizeof(float)*7);
		lastPos = seed.pos;
		lastDir = seed.dir;
		pB[0] = lastPos;
		for (bnum = 1; bnum < maxSteps; ++bnum) {
			CPoint3F nextPos;
			CVector3F nextDir;

			bool flag = true;
			for (int step = 0; step < innerSteps; ++step) {
				nextPos = lastPos - lastDir * stepSize;

				if (GetNextDirection(nextPos, &lastDir, lastDTI, &nextDir) == 0) {
					flag = false;
					break;
				}
				lastDir = nextDir;
				lastPos = nextPos;
			}
			if (flag == false)
				break;

			pB[bnum] = nextPos;
			lastDir = nextDir;
			lastPos = nextPos;
		}

		(*fibers)[i].CreateFiber(pF, pB, fnum, bnum, &seed);

		SafeDeleteArray(pF);
		SafeDeleteArray(pB);
	}

	return count;
}

void TwoTensorFiberTracking()
{
	uint w = pDTI->GetWidth();
	uint h = pDTI->GetHeight();
	uint d = pDTI->GetDepth();

	int sw = (int)w * scale;
	int sh = (int)h * scale;
	int sd = (int)d * scale;

	FILE *fp = fopen(pFiberPathname, "w");
	fprintf(fp, "%d %d %d %d %f %f\n", sw, sh, sd, scale, stepSize, maxLength);

	std::string pathname = pFiberPathname;
	size_t pos = pathname.find_last_of(SLASH);
	std::string path = pathname.substr(0, pos+1);

	for (int y = 128; y < 129; ++y) {
		printf("Fiber tracking on slice y = %d ... ", y);

		CFiber **fibers = new CFiber*[sw*sd];
		int *count = new int[sw*sd];

		for (int z = 0; z < sd; ++z) {
		//omp_set_num_threads(16);
//#pragma omp parallel for
			for (int x = 0; x < sw; ++x) {
				CPoint3F pos = CPoint3F();
				pos.m_x = (float)x/(float)scale + 1.0f / ((float)scale * 2.0f);
				pos.m_y = (float)y/(float)scale + 1.0f / ((float)scale * 2.0f);
				pos.m_z = (float)z/(float)scale + 1.0f / ((float)scale * 2.0f);
				int index = z * sw + x;
				count[index] = TrackingFibers(&(fibers[index]), pos);
			}
			float ff = float(z+1) / float(sd) * 100.0f;
			printf("\rFiber tracking on slice y = %d ... %.3f%%", y, ff);
		}

		char confPathname[PATHNAME_SIZE];
		char vertPathname[PATHNAME_SIZE];
		sprintf(confPathname, "%sslice_y=%08.04f.conf", path.c_str(), (float)y/(float)scale);
		sprintf(vertPathname, "%sslice_y=%08.04f.vert", path.c_str(), (float)y/(float)scale);
		fprintf(fp, "%s\t%s\n", confPathname, vertPathname);

		//SaveFibers(confPathname, vertPathname, fibers, count, sw, sd);
		SaveFibersAsVTK("C:\\Ziang\\Dataset\\Alex\\2X\\ForFiberTracking\\test.vtk", fibers, count, sw, sd);

		for (int i = 0; i < sd*sw; ++i) {
			SafeDeleteArray(fibers[i]);
		}
		SafeDeleteArray(fibers);
		SafeDeleteArray(count);

		printf("\rFiber tracking on slice y = %d ... 100.0%% Done.\n", y);
	}

	fclose(fp);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pDTI = new CDTI();
	pDTI->OpenDTIFile(pDTIPathname);
	pDWI = new CDWI();
	pDWI->OpenDWIFile(pDWIPathname);
	pWMMASK = new CScalar();
	pWMMASK->OpenScalarFile(pWMMASKPathname);

	CreateMLSSamplePoints();

	TwoTensorFiberTracking();

	SafeDelete(pDTI);
	SafeDelete(pDWI);
	SafeDelete(pWMMASK);

	return 0;
}